#!/usr/bin/env python3
"""
教育LLM项目性能测试脚本
使用asyncio和aiohttp实现高并发测试
"""

import asyncio
import aiohttp
import time
import json
import random
from typing import List, Dict, Any
from dataclasses import dataclass
from datetime import datetime

@dataclass
class TestResult:
    """测试结果数据类"""
    total_requests: int
    successful_requests: int
    failed_requests: int
    total_time: float
    avg_response_time: float
    min_response_time: float
    max_response_time: float
    rps: float  # Requests Per Second
    error_rate: float

class EducationLLMPerformanceTest:
    """教育LLM项目性能测试类"""
    
    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url
        self.session = None
        self.test_data = self._generate_test_data()
        
    def _generate_test_data(self) -> Dict[str, Any]:
        """生成测试数据"""
        questions = [
            "什么是人工智能？",
            "机器学习和深度学习有什么区别？",
            "如何理解神经网络？",
            "自然语言处理的应用有哪些？",
            "计算机视觉的发展历程是什么？",
            "什么是强化学习？",
            "大数据技术有哪些核心组件？",
            "云计算的服务模式有哪些？",
            "区块链技术的特点是什么？",
            "物联网的关键技术有哪些？"
        ]
        
        subjects = ["数学", "物理", "化学", "生物", "历史", "地理", "语文", "英语", "政治", "计算机"]
        
        return {
            "questions": questions,
            "subjects": subjects,
            "users": [f"user_{i}" for i in range(1, 101)]  # 100个测试用户
        }
    
    async def _create_session(self):
        """创建aiohttp会话"""
        if self.session is None:
            self.session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=30),
                connector=aiohttp.TCPConnector(limit=1000, limit_per_host=100)
            )
    
    async def _close_session(self):
        """关闭aiohttp会话"""
        if self.session:
            await self.session.close()
            self.session = None
    
    async def _test_query_api(self, user_id: str, question: str) -> Dict[str, Any]:
        """测试知识问答API"""
        url = f"{self.base_url}/api/query"
        payload = {
            "question": question,
            "user_id": user_id
        }
        
        start_time = time.time()
        try:
            async with self.session.post(url, json=payload) as response:
                response_time = time.time() - start_time
                response_data = await response.json() if response.content_length else {}
                return {
                    "success": response.status == 200,
                    "status_code": response.status,
                    "response_time": response_time,
                    "response_data": response_data
                }
        except Exception as e:
            response_time = time.time() - start_time
            return {
                "success": False,
                "status_code": 0,
                "response_time": response_time,
                "error": str(e)
            }
    
    async def _test_recommendation_api(self, user_id: str, subject: str) -> Dict[str, Any]:
        """测试学习路径推荐API"""
        url = f"{self.base_url}/api/recommend"
        payload = {
            "user_id": user_id,
            "subject": subject
        }
        
        start_time = time.time()
        try:
            async with self.session.post(url, json=payload) as response:
                response_time = time.time() - start_time
                response_data = await response.json() if response.content_length else {}
                return {
                    "success": response.status == 200,
                    "status_code": response.status,
                    "response_time": response_time,
                    "response_data": response_data
                }
        except Exception as e:
            response_time = time.time() - start_time
            return {
                "success": False,
                "status_code": 0,
                "response_time": response_time,
                "error": str(e)
            }
    
    async def _test_user_profile_api(self, user_id: str) -> Dict[str, Any]:
        """测试用户画像API"""
        url = f"{self.base_url}/api/user_profile/{user_id}"
        
        start_time = time.time()
        try:
            async with self.session.get(url) as response:
                response_time = time.time() - start_time
                response_data = await response.json() if response.content_length else {}
                return {
                    "success": response.status == 200,
                    "status_code": response.status,
                    "response_time": response_time,
                    "response_data": response_data
                }
        except Exception as e:
            response_time = time.time() - start_time
            return {
                "success": False,
                "status_code": 0,
                "response_time": response_time,
                "error": str(e)
            }
    
    async def run_concurrent_test(self, 
                                  concurrent_users: int, 
                                  requests_per_user: int,
                                  test_duration: int = None) -> TestResult:
        """
        运行并发测试
        
        Args:
            concurrent_users: 并发用户数
            requests_per_user: 每个用户的请求数
            test_duration: 测试持续时间（秒），如果指定则忽略requests_per_user
        """
        await self._create_session()
        
        print(f"开始并发测试: {concurrent_users} 并发用户")
        
        start_time = time.time()
        all_results = []
        
        # 创建测试任务
        tasks = []
        for i in range(concurrent_users):
            user_id = random.choice(self.test_data["users"])
            task = self._run_user_session(user_id, requests_per_user, test_duration)
            tasks.append(task)
        
        # 执行所有测试任务
        user_results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 合并结果
        for result in user_results:
            if isinstance(result, Exception):
                print(f"用户会话出错: {result}")
                continue
            all_results.extend(result)
        
        total_time = time.time() - start_time
        
        # 计算统计信息
        successful_requests = sum(1 for r in all_results if r.get("success", False))
        failed_requests = len(all_results) - successful_requests
        
        response_times = [r["response_time"] for r in all_results if r.get("success", False)]
        
        if response_times:
            avg_response_time = sum(response_times) / len(response_times)
            min_response_time = min(response_times)
            max_response_time = max(response_times)
        else:
            avg_response_time = min_response_time = max_response_time = 0
        
        total_requests = len(all_results)
        rps = total_requests / total_time if total_time > 0 else 0
        error_rate = failed_requests / total_requests if total_requests > 0 else 0
        
        test_result = TestResult(
            total_requests=total_requests,
            successful_requests=successful_requests,
            failed_requests=failed_requests,
            total_time=total_time,
            avg_response_time=avg_response_time,
            min_response_time=min_response_time,
            max_response_time=max_response_time,
            rps=rps,
            error_rate=error_rate
        )
        
        return test_result
    
    async def _run_user_session(self, 
                                user_id: str, 
                                requests_count: int, 
                                test_duration: int = None) -> List[Dict]:
        """运行单个用户会话"""
        results = []
        start_time = time.time()
        
        request_count = 0
        while True:
            # 检查是否达到请求次数或时间限制
            if test_duration:
                if time.time() - start_time > test_duration:
                    break
            else:
                if request_count >= requests_count:
                    break
            
            # 随机选择API进行测试
            api_choice = random.choice(["query", "recommend", "profile"])
            
            if api_choice == "query":
                question = random.choice(self.test_data["questions"])
                result = await self._test_query_api(user_id, question)
            elif api_choice == "recommend":
                subject = random.choice(self.test_data["subjects"])
                result = await self._test_recommendation_api(user_id, subject)
            else:  # profile
                result = await self._test_user_profile_api(user_id)
            
            results.append(result)
            request_count += 1
            
            # 添加小延迟以避免过于激进的请求
            await asyncio.sleep(random.uniform(0.01, 0.1))
        
        return results
    
    def print_test_result(self, result: TestResult):
        """打印测试结果"""
        print("\n" + "="*60)
        print("性能测试结果")
        print("="*60)
        print(f"总请求数: {result.total_requests}")
        print(f"成功请求数: {result.successful_requests}")
        print(f"失败请求数: {result.failed_requests}")
        print(f"错误率: {result.error_rate:.2%}")
        print(f"总耗时: {result.total_time:.2f} 秒")
        print(f"平均响应时间: {result.avg_response_time*1000:.2f} 毫秒")
        print(f"最小响应时间: {result.min_response_time*1000:.2f} 毫秒")
        print(f"最大响应时间: {result.max_response_time*1000:.2f} 毫秒")
        print(f"请求速率: {result.rps:.2f} RPS")
        print("="*60)

async def main():
    """主函数"""
    # 创建测试实例
    tester = EducationLLMPerformanceTest("http://localhost:8000")
    
    try:
        # 测试阶段1: 低并发测试 (100并发用户，每个用户5个请求)
        print("阶段1: 低并发测试 (100并发用户)")
        result1 = await tester.run_concurrent_test(concurrent_users=100, requests_per_user=5)
        tester.print_test_result(result1)
        
        # 等待一段时间让系统恢复
        await asyncio.sleep(5)
        
        # 测试阶段2: 中并发测试 (300并发用户，每个用户3个请求)
        print("\n阶段2: 中并发测试 (300并发用户)")
        result2 = await tester.run_concurrent_test(concurrent_users=300, requests_per_user=3)
        tester.print_test_result(result2)
        
        # 等待一段时间让系统恢复
        await asyncio.sleep(5)
        
        # 测试阶段3: 高并发测试 (500并发用户，每个用户2个请求)
        print("\n阶段3: 高并发测试 (500并发用户)")
        result3 = await tester.run_concurrent_test(concurrent_users=500, requests_per_user=2)
        tester.print_test_result(result3)
        
    except KeyboardInterrupt:
        print("\n测试被用户中断")
    except Exception as e:
        print(f"测试过程中发生错误: {e}")
    finally:
        await tester._close_session()

if __name__ == "__main__":
    asyncio.run(main())